
import numpy as np
import torch
def window_partition(x, window_size):
    """
    Args:
        x: (B, H, W, C)
        window_size (int): window size

    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B,C, H, W = x.shape
    x = x.view(B,C, H // window_size, window_size, W // window_size, window_size)
    x=x.permute(0, 2, 4, 1, 3, 5) #B, H // window_size, W // window_size,C, window_size, window_size

    windows = x.reshape(-1,C,window_size, window_size) # -1,C,winH,winW
    return windows
def window_reverse(windows, H, W):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image

    Returns:
        x: (B, H, W, C)
    """
    _,C,window_size,_=windows.shape

    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(B, H // window_size, W // window_size, C, window_size, window_size)
    x= x.permute(0,3, 1, 4, 2, 5)#B, C, H // window_size, window_size, W // window_size, window_size
    x=x.reshape(B,C,H,W)
    return x

x=[
    [
        [
            [[1],[2],[3],[4]],[[4],[5],[6],[7]],[[7],[8],[9],[10]],[[10],[11],[12],[13]]
        ],
        [
            [[51],[52],[53],[54]],[[54],[55],[56],[57]],[[57],[58],[59],[510]],[[510],[511],[512],[513]]
        ]
    ]
]
x=np.random.random((1,3,4,4))
x=torch.from_numpy(x)
print(x)
x=window_partition(x,2)
print(x.shape)
x=window_reverse(x,4,4)
print(x)